import datetime
import logging
import time

import asyncio
from fastapi import FastAPI
from starlette.middleware.gzip import GZipMiddleware
from starlette.requests import Request


class RequestRecorder:
    records = {}
    lock = asyncio.Lock()

    @staticmethod
    async def update(
            url: str,
            method: str,
            delay: float,
    ):
        async with RequestRecorder.lock:
            if (method, url) not in RequestRecorder.records:
                RequestRecorder.records[(method, url)] = {
                    "max_delay": delay,
                    "min_delay": delay,
                    "avg_delay": delay,
                    "cnt": 1,

                }
                return

            record = RequestRecorder.records[(method, url)]
            record["max_delay"] = max(record["max_delay"], delay)
            record["min_delay"] = min(record["min_delay"], delay)
            record["avg_delay"] = (record["avg_delay"] * record["cnt"] + delay) / (record["cnt"] + 1)
            record["cnt"] += 1


def get_stats():
    stats = []
    for (method, url), info in RequestRecorder.records.items():
        stats.append({
            "method": method,
            "url": url,
            "info": info,
        })
    return stats


def init_middleware(app: FastAPI, enable_middleware: bool = False):
    app.add_middleware(GZipMiddleware, minimum_size=1000, compresslevel=5)

    if not enable_middleware:
        return

    @app.middleware("http")
    async def latency_detect(request: Request, call_next):
        url = request.url.path
        method = request.method

        request_arrive_time = int(time.time() * 1000)
        request_start_time = request.headers.get("X-Request-Start-Time")
        response = await call_next(request)
        response_start_time = int(time.time() * 1000)
        response.headers["X-Response-Start-Time"] = str(response_start_time)
        response.headers["X-Request-Arrive-Time"] = str(request_arrive_time)
        if request_start_time:
            response.headers["X-Client-Delay"] = str(request_arrive_time - int(request_start_time))

        print("[Latency Middleware] {} {} {} Request took {:.4f} ms]".format(
            datetime.datetime.now(),
            method,
            url,
            response_start_time - request_arrive_time))

        await RequestRecorder.update(url, method, delay=response_start_time - request_arrive_time)
        return response
